package com.huangkailong.learning.datastructure.tree;

import com.huangkailong.learning.datastructure.common.Assert;
import com.huangkailong.learning.datastructure.tree.printer.BinaryTreeInfo;
import com.huangkailong.learning.datastructure.tree.printer.BinaryTrees;

import java.util.Objects;

/**
 * 红黑树
 *
 * @author huangkl
 * @since 1.0.0
 */
public class RedBlackTree<K extends Comparable<K>, V> extends AbstractBinaryFindTree<K,V>{
    private static final boolean RED = true;
    private static final boolean BLACK = false;
    private Node<K, V> root;
    private int size;

    protected static class Node<K extends Comparable<K>, V> extends AbstractBinaryFindTree.Node<K,V>{
        /**
         * 父结点指针
         */
        Node<K, V> parent;

        /**
         * 左孩子结点指针
         */
        Node<K, V> left;

        /**
         * 右孩子结点指针
         */
        Node<K, V> right;

        /**
         * 颜色, ture 红色， false 黑色
         */
        boolean color;

        public Node(K key, V value) {
            super(key, value);
        }

        public Node(K key, V value, boolean color) {
            this(key, value);
            this.color = color;
        }

        public Node(K key, V value, Node<K, V> parent) {
            this(key, value);
            this.parent = parent;
        }

        @Override
        AbstractBinaryFindTree.Node<K, V> getLeft() {
            return left;
        }

        @Override
        AbstractBinaryFindTree.Node<K, V> getRight() {
            return right;
        }
    }

    @Override
    protected AbstractBinaryFindTree.Node<K, V> getRoot() {
        return root;
    }

    /**
     * 沿着 node 左旋
     *
     * @param node {@link Node}
     */
    private void rotateLeft(Node<K, V> node){
        final Node<K, V> parent = parentOf(node);
        final Node<K, V> right = rightOf(node);

        node.right = right.left;
        if(Objects.nonNull(node.right)){
            node.right.parent = node;
        }
        right.left = node;
        changeParentChild(parent, node, right);
        right.parent = parent;
        node.parent = right;

    }

    /**
     * 沿着 node 右旋
     *
     * @param node {@link Node}
     */
    private void rotateRight(Node<K, V> node){
        final Node<K, V> parent = parentOf(node);
        final Node<K, V> left = leftOf(node);

       node.left = left.right;
       if(Objects.nonNull(node.left)){
           node.left.parent = node;
       }
       left.right = node;
       changeParentChild(parent, node, left);
       left.parent = parent;
       node.parent = left;
    }


    private void changeParentChild(Node<K, V> parent, Node<K, V> oldChild, Node<K, V> newChild){
        if(Objects.isNull(parent)){
            // oldChild 是根结点
            this.root = newChild;
            return;
        }

        if(parent.left == oldChild){
            // oldChild 是 parent 的左孩子
            parent.left = newChild;
        }else {
            // oldChild 是 parent 的右孩子
            parent.right = newChild;
        }
    }

    @Override
    public V put(K key, V value){
        size++;

        // 根结点直接插入
        if(Objects.isNull(this.root)){
            this.root = new Node<>(key, value, BLACK);
            return null;
        }

        // 找到插入结点, pp 指向 p 的父结点
        Node<K, V> p = this.root;
        Node<K, V> pp = null;
        while (Objects.nonNull(p)){
            pp = p;
            if(lt(key, p.key)){
                p = p.left;
            }else if(gt(key, p.key)){
                p = p.right;
            }else {
                p.key = key;
                p.value = value;
                return null;
            }
        }

        // 插入结点
        Node<K, V> node = insertNode(key, value, pp);

        // 结点调整
        fixAfterPut(node);

        return null;
    }

    @Override
    public V delete(K key) {
        Assert.isNotNull(key, "Key 不能为空!");
        Node<K, V> root = this.root;
        if(Objects.isNull(root)){
            return null;
        }

        Node<K, V> node = findNode(root, key);
        if(Objects.isNull(node)){
            return null;
        }
        doDelete(node);
        size--;
        return node.value;
    }

    private Node<K, V> findNode(Node<K,V> root, K key){
        Node<K, V> p = root;
        while (Objects.nonNull(p)){
            if(eq(key, p.key)){
                return p;
            }else if (lt(key, p.key)){
                p = p.left;
            }else {
                p = p.right;
            }
        }
        return null;
    }

    private void doDelete(Node<K,V> node){
        Node<K, V> p = node;
        if(Objects.nonNull(leftOf(node)) && Objects.nonNull(rightOf(node))){
            // 有左孩子和右孩子
            Node<K, V> predecessor = getPredecessor(node);
            node.key = predecessor.key;
            node.value = predecessor.value;
            p = predecessor;
        }

        Node<K, V> replaceNode = p == null ? null : p.left == null ? p.right : p.left;
        if(Objects.nonNull(replaceNode)){
            // 有左孩子或右孩子
            replaceNode.parent = p.parent;
            if(Objects.nonNull(p.parent)){
                if(leftOf(p.parent) == p){
                    p.parent.left = replaceNode;
                }else {
                    p.parent.right = replaceNode;
                }
            }else {
                this.root = replaceNode;
            }

            if(colorIs(p, BLACK)){
                fixAfterDelete(replaceNode);
            }
        }else {
            // 无左孩子或右孩子
            if(colorIs(p, BLACK)){
                fixAfterDelete(p);
            }
            if(Objects.nonNull(p.parent)){
                if(leftOf(p.parent) == p){
                    p.parent.left = null;
                }else {
                    p.parent.right = null;
                }
            }else {
                this.root = null;
            }
        }
    }

    private void fixAfterDelete(Node<K, V> node){
        while (node != root && node != null && colorIs(node, BLACK)){
            Node<K, V> nodeParent = parentOf(node);
            if(leftOf(nodeParent) == node){
                // 左孩子
                // nodeBrother 表示的是 node 的兄弟结点，此兄弟结点指的是234树的兄弟结点而不是二叉树
                Node<K, V> nodeBrother = rightOf(nodeParent);
                if(colorIs(nodeBrother, RED)){
                    // nodeBrother 为红色，表示还不是真正的兄弟结点，所以需要进过变色和旋转
                    /**
                     *  ┌─2─┐
                     *  │   │
                     *  1 ┌─^6─┐
                     *    │    │
                     *  ┌─4─┐  7
                     *  │   │
                     * ^3   ^5
                     *
                     * node --> 1
                     * nodeParent --> 2
                     * nodeBrother --> 6
                     */
                    setColor(nodeBrother, BLACK);
                    setColor(nodeParent, RED);
                    rotateLeft(nodeParent);
                    // 指向了真正的兄弟结点
                    nodeBrother = rightOf(nodeParent);
                    /**
                     *    ┌─6─┐
                     *    │   │
                     * ┌─^2─┐ 7
                     * │    │
                     * 1  ┌─4─┐
                     *    │   │
                     *   ^3   ^5
                     * node --> 1
                     * nodeParent --> 2
                     * nodeBrother --> 4
                     */
                }
                if(colorIs(leftOf(nodeBrother), BLACK) && colorIs(rightOf(nodeBrother), BLACK)){
                    // 无左右孩子结点或左右孩子都是黑色结点
                    setColor(nodeBrother, RED);
                    node = nodeParent;
                }else {
                    if(rightOf(nodeBrother) == null){
                        // 兄弟结点无右孩子
                        /**
                         * ┌─2─┐
                         * │   │
                         * 1 ┌─4
                         *   │
                         *  ^3
                         * node --> 1
                         * nodeParent --> 2
                         * nodeBrother --> 4
                         */
                        setColor(leftOf(nodeBrother), BLACK);
                        setColor(nodeBrother, RED);
                        rotateRight(nodeBrother);
                        nodeBrother = rightOf(nodeParent);
                        /**
                         * ┌─2─┐
                         * │   │
                         * 1   3─┐
                         *       │
                         *       ^4
                         * node --> 1
                         * nodeParent --> 2
                         * nodeBrother --> 3
                         */
                    }
                    setColor(nodeBrother, colorOf(nodeParent));
                    setColor(nodeParent, BLACK);
                    setColor(rightOf(nodeBrother), BLACK);
                    rotateLeft(nodeParent);
                    return;
                }

            }else {
                // 右孩子
                Node<K, V> nodeBrother = leftOf(nodeParent);
                if(colorIs(nodeBrother, RED)){
                    // nodeBrother 为红色，表示还不是真正的兄弟结点，所以需要进过变色和旋转
                    setColor(nodeBrother, BLACK);
                    setColor(nodeParent, RED);
                    rotateRight(nodeParent);
                    // 指向了真正的兄弟结点
                    nodeBrother = leftOf(nodeParent);
                }
                if(colorIs(leftOf(nodeBrother), BLACK) && colorIs(rightOf(nodeBrother), BLACK)){
                    // 无左右孩子结点或左右孩子结点都是黑色结点
                    setColor(nodeBrother, RED);
                    node = nodeParent;
                }else {
                    if(leftOf(nodeBrother) == null){
                        // 兄弟结点无右孩子
                        setColor(leftOf(nodeBrother), BLACK);
                        setColor(nodeBrother, RED);
                        rotateLeft(nodeBrother);
                        nodeBrother = leftOf(nodeParent);
                    }
                    setColor(nodeBrother, colorOf(nodeParent));
                    setColor(nodeParent, BLACK);
                    setColor(leftOf(nodeBrother), BLACK);
                    rotateRight(nodeParent);
                    return;
                }
            }
        }
        setColor(node, BLACK);
    }

    /**
     * 获取 node 结点的前继结点
     *
     * @param node {@link Node}
     * @return node 结点的前继结点
     */
    private Node<K, V> getPredecessor(Node<K,V> node){
        Node<K, V> p = node.left;
        Node<K, V> pp = p;
        while (Objects.nonNull(p)){
            pp = p;
            p = p.right;
        }
        return pp;
    }

    /**
     * 获取 node 结点的后继结点
     *
     * @param node {@link Node}
     * @return node 结点的后继结点
     */
    private Node<K, V> getSuccessor(Node<K,V> node){
        Node<K, V> p = node.right;
        Node<K, V> pp = p;
        while (Objects.nonNull(p)){
            pp = p;
            p = p.left;
        }
        return pp;
    }

    @Override
    public int size() {
        return size;
    }

    /**
     * 插入结点
     *
     * @param key key
     * @param value value
     * @param parentNodeOfInsert 插入结点的父结点
     * @return 插入的结点
     */
    private Node<K, V> insertNode(K key, V value, Node<K, V> parentNodeOfInsert) {
        Node<K, V> node = new Node<>(key, value, parentNodeOfInsert);
        if(lt(key, parentNodeOfInsert.key)){
            parentNodeOfInsert.left = node;
        }else {
            parentNodeOfInsert.right = node;
        }
        return node;
    }

    private void fixAfterPut(Node<K, V> node) {
        // 插入的结点是红色
        setColor(node, RED);

        // p 父亲结点为红色结点才进行调整
        Node<K, V> p = node;
        while (p != this.root && parentOf(p) != null && colorOf(parentOf(p)) == RED){
            // 父结点
            Node<K, V> parent = parentOf(p);
            // 爷爷结点
            Node<K, V> grandpa = parentOf(parent);
            // 父结点的兄弟结点
            Node<K, V> parentBrother;
            if(leftOf(grandpa) == parent){
                parentBrother = rightOf(grandpa);
            }else {
                parentBrother = leftOf(grandpa);
            }

            // 根据情况进行处理
            if(colorIs(parentBrother, BLACK)){
                // 父亲兄弟结点为黑色
                if(leftOf(parent) == p && leftOf(grandpa) == parent){
                    // 插入结点是父结点的左孩子，插入结点的父亲结点是插入结点的爷爷结点的左孩子
                    rotateRight(grandpa);
                    setColor(parent, BLACK);
                    setColor(grandpa, RED);
                }else if(leftOf(parent) == p && rightOf(grandpa) == parent){
                    // 插入结点是父结点的左孩子，插入结点的父亲结点是插入结点的爷爷结点的右孩子
                    rotateRight(parent);
                    rotateLeft(grandpa);
                    setColor(grandpa, RED);
                    setColor(p, BLACK);
                }else if(rightOf(parent) == p && rightOf(grandpa) == parent){
                    // 插入结点是父结点的右孩子，插入结点的父亲结点是插入结点的爷爷结点的右孩子
                    rotateLeft(grandpa);
                    setColor(grandpa, RED);
                    setColor(parent, BLACK);
                }else if(rightOf(parent) == p && leftOf(grandpa) == parent){
                    // 插入结点是父结点的右孩子，插入结点的父亲结点是插入结点的爷爷结点的左孩子
                    rotateLeft(parent);
                    rotateRight(grandpa);
                    setColor(grandpa, RED);
                    setColor(p, BLACK);
                }
                break;
            }else {
                // 父亲兄弟结点为红色
                setColor(parent, BLACK);
                setColor(parentBrother, BLACK);
                setColor(grandpa, RED);
                p = grandpa;
            }
        }

        // 根结点永远为黑色
        setColor(this.root, BLACK);
    }

    private Node<K, V> parentOf(Node<K, V> node){
        return node == null ? null : node.parent;
    }

    private Node<K, V> leftOf(Node<K, V> node){
        return node == null ? null : node.left;
    }

    private Node<K, V> rightOf(Node<K, V> node){
        return node == null ? null : node.right;
    }

    private boolean colorOf(Node<K, V> node){
        return node == null ? BLACK : node.color;
    }

    private void setColor(Node<K, V> node, boolean color){
        if(node != null){
            node.color = color;
        }
    }

    /**
     * 判断 node 的 color 是否为 color
     *
     * @param node {@link Node}
     * @param color ture 为红色， false 为黑色
     * @return node.color == color
     */
    private boolean colorIs(Node<K, V> node, boolean color){
        return colorOf(node) == color;
    }

    @Override
    public String printTree(BinaryTrees.PrintStyle printStyle) {
        return BinaryTrees.printString(new NodeInfo(this.root), printStyle);
    }

    private class NodeInfo implements BinaryTreeInfo {
        private Node root;

        public NodeInfo(Node root) {
            this.root = root;
        }

        @Override
        public Object root() {
            return root;
        }

        @Override
        public Object left(Object node) {
            return ((Node) node).left;
        }

        @Override
        public Object right(Object node) {
            return ((Node) node).right;
        }

        @Override
        public Object string(Object node) {
            Node n = (Node) node;
            if(colorIs(n, BLACK)){
                return n.key;
            }

            return "^" + n.key;
        }
    }
}
